package com.ksyun.campus.metaserver.services;

import com.ksyun.campus.metaserver.domain.DataServerInfo;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.stream.Collectors;

/**
 * 副本分配监控器
 * 监控三副本分配的效果，确保负载均衡和数据安全性
 */
@Slf4j
@Component
public class ReplicaAllocationMonitor {
    
    // 统计每个DataServer被选择的次数
    private final Map<String, AtomicLong> dataServerSelectionCount = new ConcurrentHashMap<>();
    
    // 统计每个DataServer作为PRIMARY副本的次数
    private final Map<String, AtomicLong> primaryReplicaCount = new ConcurrentHashMap<>();
    
    // 记录最近的分配历史
    private final List<AllocationRecord> recentAllocations = Collections.synchronizedList(new ArrayList<>());
    private static final int MAX_RECENT_RECORDS = 1000;
    
    // 总分配统计
    private final AtomicLong totalBlocks = new AtomicLong(0);
    private final AtomicLong totalReplicas = new AtomicLong(0);
    
    /**
     * 记录一次块分配
     */
    public void recordAllocation(String fileId, int blkIdx, List<DataServerInfo> selectedServers) {
        if (selectedServers.isEmpty()) {
            log.warn("Empty server list for block {}:{}", fileId, blkIdx);
            return;
        }
        
        totalBlocks.incrementAndGet();
        totalReplicas.addAndGet(selectedServers.size());
        
        // 记录每个DataServer的选择次数
        for (int i = 0; i < selectedServers.size(); i++) {
            DataServerInfo server = selectedServers.get(i);
            String dsId = server.getDsId();
            
            // 增加选择计数
            dataServerSelectionCount.computeIfAbsent(dsId, k -> new AtomicLong(0))
                    .incrementAndGet();
            
            // 如果是第一个副本（PRIMARY），增加PRIMARY计数
            if (i == 0) {
                primaryReplicaCount.computeIfAbsent(dsId, k -> new AtomicLong(0))
                        .incrementAndGet();
            }
        }
        
        // 记录分配历史
        AllocationRecord record = new AllocationRecord();
        record.fileId = fileId;
        record.blkIdx = blkIdx;
        record.timestamp = LocalDateTime.now();
        record.serverIds = selectedServers.stream()
                .map(DataServerInfo::getDsId)
                .collect(Collectors.toList());
        record.primaryServerId = selectedServers.get(0).getDsId();
        
        synchronized (recentAllocations) {
            recentAllocations.add(record);
            if (recentAllocations.size() > MAX_RECENT_RECORDS) {
                recentAllocations.remove(0);
            }
        }
        
        // 定期输出统计信息
        if (totalBlocks.get() % 50 == 0) {
            logStatistics();
        }
    }
    
    /**
     * 获取负载均衡统计信息
     */
    public LoadBalanceStats getLoadBalanceStats() {
        LoadBalanceStats stats = new LoadBalanceStats();
        
        if (dataServerSelectionCount.isEmpty()) {
            return stats;
        }
        
        // 计算DataServer选择分布
        Collection<AtomicLong> selections = dataServerSelectionCount.values();
        long totalSelections = selections.stream().mapToLong(AtomicLong::get).sum();
        
        if (totalSelections > 0) {
            double avgSelections = (double) totalSelections / selections.size();
            
            // 计算标准差
            double variance = selections.stream()
                    .mapToDouble(count -> Math.pow(count.get() - avgSelections, 2))
                    .average()
                    .orElse(0.0);
            stats.selectionStdDev = Math.sqrt(variance);
            stats.avgSelectionsPerServer = avgSelections;
            
            // 计算最大最小选择比例
            long maxSelections = selections.stream().mapToLong(AtomicLong::get).max().orElse(0);
            long minSelections = selections.stream().mapToLong(AtomicLong::get).min().orElse(0);
            stats.maxMinSelectionRatio = minSelections > 0 ? (double) maxSelections / minSelections : 0.0;
        }
        
        // 计算PRIMARY副本分布
        if (!primaryReplicaCount.isEmpty()) {
            Collection<AtomicLong> primaries = primaryReplicaCount.values();
            long totalPrimaries = primaries.stream().mapToLong(AtomicLong::get).sum();
            
            if (totalPrimaries > 0) {
                double avgPrimaries = (double) totalPrimaries / primaries.size();
                double primaryVariance = primaries.stream()
                        .mapToDouble(count -> Math.pow(count.get() - avgPrimaries, 2))
                        .average()
                        .orElse(0.0);
                stats.primaryStdDev = Math.sqrt(primaryVariance);
                stats.avgPrimariesPerServer = avgPrimaries;
            }
        }
        
        return stats;
    }
    
    /**
     * 获取详细的分配统计信息
     */
    public Map<String, Object> getDetailedStats() {
        Map<String, Object> stats = new HashMap<>();
        
        // 基础统计
        stats.put("totalBlocks", totalBlocks.get());
        stats.put("totalReplicas", totalReplicas.get());
        stats.put("avgReplicasPerBlock", totalBlocks.get() > 0 ? 
                (double) totalReplicas.get() / totalBlocks.get() : 0.0);
        
        // DataServer选择分布
        Map<String, Long> selectionDistribution = dataServerSelectionCount.entrySet().stream()
                .collect(Collectors.toMap(
                        Map.Entry::getKey,
                        entry -> entry.getValue().get()
                ));
        stats.put("serverSelectionDistribution", selectionDistribution);
        
        // PRIMARY副本分布
        Map<String, Long> primaryDistribution = primaryReplicaCount.entrySet().stream()
                .collect(Collectors.toMap(
                        Map.Entry::getKey,
                        entry -> entry.getValue().get()
                ));
        stats.put("primaryReplicaDistribution", primaryDistribution);
        
        // 负载均衡统计
        stats.put("loadBalanceStats", getLoadBalanceStats());
        
        // 最近分配历史
        synchronized (recentAllocations) {
            List<Map<String, Object>> recentHistory = recentAllocations.stream()
                    .limit(10) // 只返回最近10条记录
                    .map(record -> {
                        Map<String, Object> recordMap = new HashMap<>();
                        recordMap.put("fileId", record.fileId);
                        recordMap.put("blkIdx", record.blkIdx);
                        recordMap.put("timestamp", record.timestamp);
                        recordMap.put("serverIds", record.serverIds);
                        recordMap.put("primaryServerId", record.primaryServerId);
                        return recordMap;
                    })
                    .collect(Collectors.toList());
            stats.put("recentAllocations", recentHistory);
        }
        
        return stats;
    }
    
    /**
     * 检查负载均衡是否良好
     */
    public boolean isLoadBalanced(double maxStdDevThreshold) {
        LoadBalanceStats stats = getLoadBalanceStats();
        return stats.selectionStdDev <= maxStdDevThreshold && 
               stats.primaryStdDev <= maxStdDevThreshold;
    }
    
    /**
     * 重置统计信息
     */
    public void reset() {
        dataServerSelectionCount.clear();
        primaryReplicaCount.clear();
        totalBlocks.set(0);
        totalReplicas.set(0);
        
        synchronized (recentAllocations) {
            recentAllocations.clear();
        }
        
        log.info("Replica allocation monitor statistics reset");
    }
    
    /**
     * 输出统计日志
     */
    private void logStatistics() {
        LoadBalanceStats stats = getLoadBalanceStats();
        
        log.info("=== Replica Allocation Statistics ===");
        log.info("Total blocks allocated: {}, Total replicas: {}", 
                totalBlocks.get(), totalReplicas.get());
        log.info("Average replicas per block: {:.2f}", 
                totalBlocks.get() > 0 ? (double) totalReplicas.get() / totalBlocks.get() : 0.0);
        
        log.info("Load balance - Selection StdDev: {:.2f}, Primary StdDev: {:.2f}", 
                stats.selectionStdDev, stats.primaryStdDev);
        log.info("Max/Min selection ratio: {:.2f}", stats.maxMinSelectionRatio);
        
        // 输出各DataServer的选择次数
        if (!dataServerSelectionCount.isEmpty()) {
            StringBuilder sb = new StringBuilder("DataServer selection counts: ");
            dataServerSelectionCount.entrySet().stream()
                    .sorted(Map.Entry.<String, AtomicLong>comparingByValue((a, b) -> 
                            Long.compare(b.get(), a.get())))
                    .forEach(entry -> sb.append(entry.getKey())
                            .append(":")
                            .append(entry.getValue().get())
                            .append(" "));
            log.info(sb.toString());
        }
        
        // 检查负载均衡状态
        boolean balanced = isLoadBalanced(5.0); // 标准差阈值5.0
        log.info("Load balance status: {}", balanced ? "GOOD" : "NEEDS_IMPROVEMENT");
    }
    
    // 内部数据结构
    @Data
    public static class LoadBalanceStats {
        private double selectionStdDev;
        private double avgSelectionsPerServer;
        private double maxMinSelectionRatio;
        private double primaryStdDev;
        private double avgPrimariesPerServer;
    }
    
    private static class AllocationRecord {
        String fileId;
        int blkIdx;
        LocalDateTime timestamp;
        List<String> serverIds;
        String primaryServerId;
    }
}
